import torch
import torch.nn as nn
import os
import numpy as np

from config.path_config import GET_PROJECT_ROOT
from src.model.model_point_picker import ModelPointPicker
from src.model.layers.knn_simple import KNNSimple
from src.model.layers.devconv_gnn import DevConvGNN
from src.model.layers.sparse_attention_edge_predictor_layer import SparseAttentionEdgePredictorLayer
from src.model.layers.face_candidate import FaceCandidatesLayer
from src.model.layers.triangle_indexes import TriangleIndexes
from src.model.layers.first_p_init import FirstPInitLayer
from src.model.layers.knn import KNN
from src.model.layers.r_matrix import RMatrix
from src.model.layers.mlp import MLP

PROJECT_ROOT = GET_PROJECT_ROOT()


class ModelTriangleGenerator(nn.Module):
    def __init__(self, number_neigh_tri, debug=False):
        super().__init__()
        self.debug = debug
        self.number_neigh_tri = number_neigh_tri

        k = 20
        # 定义模型的各个层
        self.model_point_picker = ModelPointPicker()
        self.model_point_picker.load_state_dict(
            torch.load(os.path.join(PROJECT_ROOT, 'save_models/point_picker/1/410760.pt')))

        self.layer_knn_simple = KNNSimple(k=15)
        self.layer_devconv_edge_predictor = DevConvGNN(3, 64, 1)
        self.layer_sparse_attention_edge_predictor = SparseAttentionEdgePredictorLayer(64)
        self.layer_face_cadidates = FaceCandidatesLayer()
        self.layer_triangle_indexes = TriangleIndexes()
        self.layer_first_p_init = FirstPInitLayer()
        self.layer_knn = KNN(k=k, batch_size=1000)
        self.layer_r_matrix = RMatrix()
        self.layer_mlp = MLP(1, 128, 1)

    def forward(self, target_number_point, original_graph):
        # POINT SAMPLER
        with torch.no_grad():
            self.score_original_points, self.generated_graph_nodes = self.model_point_picker(target_number_point,
                                                                                             original_graph)
        generated_graph_adjacency_matrix = self.layer_knn_simple(self.generated_graph_nodes)

        # EDGE PREDICTOR
        edge_index_generated_graph = generated_graph_adjacency_matrix.nonzero().T
        score_edge = self.layer_devconv_edge_predictor(edge_index_generated_graph,
                                                       self.generated_graph_nodes).squeeze().sigmoid()
        S = self.layer_sparse_attention_edge_predictor(score_edge, generated_graph_adjacency_matrix)

        # FACE CANDIDATES
        generated_probabilistic_graph_adjacency_matrix = self.layer_face_cadidates(S, generated_graph_adjacency_matrix)

        # FACE CLASSIFIER
        triangles_ids_igraph = self.layer_triangle_indexes(generated_graph_adjacency_matrix)
        triangles = self.generated_graph_nodes[triangles_ids_igraph]

        barycenters = triangles.mean(dim=1)
        self.original_barycenters = barycenters

        indices_neigh_tri = self.layer_knn(barycenters).int()

        r_matrix = self.layer_r_matrix(triangles, barycenters, indices_neigh_tri, self.number_neigh_tri)

        p_init = self.layer_first_p_init(triangles_ids_igraph, generated_probabilistic_graph_adjacency_matrix,
                                         triangles)
        self.final_scores = self.layer_mlp(p_init, r_matrix, indices_neigh_tri)

        target_number_triangles = target_number_point // 3
        self.selected_triangles_probabilities, self.selected_triangles_indexes = torch.topk(self.final_scores,
                                                                                            k=target_number_triangles)
        selected_triangles = triangles[self.selected_triangles_indexes]

        # print(f"type of selected_triangles: {type(selected_triangles)}")
        # print(f"type of selected_triangles: {selected_triangles}")
        # # 后处理：确保网格封闭
        # selected_triangles = self.ensure_mesh_closure(selected_triangles)
        # print(f"type of selected_triangles: {type(selected_triangles)}")
        # print(f"type of selected_triangles: {selected_triangles}")
        return selected_triangles

    def ensure_mesh_closure(self, selected_triangles):
        """
        确保生成的三角形网格封闭。通过检查边的连接性，确保网格是连通的。
        """
        # 1. 获取所有顶点和三角形
        device = selected_triangles.device  # 保持设备一致
        vertices = selected_triangles.reshape(-1, 3)  # 不用调用 cpu()，保持原设备
        # edges = self.get_edges_from_triangles(selected_triangles)

        # 2. 合并重复顶点
        unique_vertices, unique_indices = np.unique(vertices.cpu().numpy(), axis=0, return_inverse=True)


        # 3. 创建面（face）索引，并使用去重后的顶点索引
        faces = unique_indices.reshape(-1, 3)

        # 4. 构建一个顶点连接图，确保网格中的顶点之间能够正确连接
        vertex_connections = self.build_vertex_connection_graph(faces)

        # 5. 修复孤立顶点：检查未连接的顶点并添加连接
        faces = self.fix_isolated_vertices(faces, vertex_connections)

        # 6. 修复未连接的边
        faces = self.fix_unconnected_edges(faces, vertex_connections)

        # 7. 返回修复后的三角形面
        # 如果需要返回合并后的数据：
        selected_triangles = torch.tensor(faces)  # 你可以将 faces 重新构建为需要的张量格式
        return selected_triangles

    def get_edges_from_triangles(self, triangles):
        """
        从三角形中提取边的信息
        """
        edges = []
        for triangle in triangles:
            # 确保triangle的元素是Tensor类型并且形状一致
            edge1 = torch.tensor([triangle[0], triangle[1]], dtype=torch.long, )
            edge2 = torch.tensor([triangle[1], triangle[2]], dtype=torch.long, )
            edge3 = torch.tensor([triangle[2], triangle[0]], dtype=torch.long, )

            edges.append(edge1)
            edges.append(edge2)
            edges.append(edge3)

        # 在堆叠之前，检查所有的边是否都是正确的Tensor并且形状一致
        try:
            stacked_edges = torch.stack(edges)
        except Exception as e:
            print(f"Error while stacking edges: {e}")
            print(f"edges: {edges}")
            raise e

        return stacked_edges

    def build_vertex_connection_graph(self, faces):
        """
        构建一个顶点连接图，确保网格中的顶点之间能够正确连接
        """
        vertex_connections = {}
        for face in faces:
            for i in range(3):
                if face[i] not in vertex_connections:
                    vertex_connections[face[i]] = set()
                vertex_connections[face[i]].add(face[(i + 1) % 3])
        return vertex_connections

    def fix_isolated_vertices(self, faces, vertex_connections):
        """
        修复孤立顶点：检查未连接的顶点并添加连接
        """
        all_vertices = set(np.unique(faces))  # 获取所有顶点
        connected_vertices = set(vertex_connections.keys())  # 获取已连接的顶点

        isolated_vertices = all_vertices - connected_vertices  # 找出孤立顶点
        for isolated_vertex in isolated_vertices:
            # 为孤立顶点找到最接近的其他顶点并进行连接
            closest_vertex = self.find_closest_connected_vertex(isolated_vertex, vertex_connections)
            # 将孤立顶点与最近的顶点连接
            faces = self.add_connection(faces, isolated_vertex, closest_vertex)
        return faces

    def find_closest_connected_vertex(self, isolated_vertex, vertex_connections):
        """
        找到与孤立顶点最接近的已连接顶点
        """
        # 在此假设每个顶点都可以通过欧几里得距离找到最近的连接点
        # 这里的实现简化为返回连接图中的一个任意点
        connected_vertices = list(vertex_connections.keys())
        return connected_vertices[0]

    def add_connection(self, faces, isolated_vertex, connected_vertex):
        """
        在三角形面中添加连接
        """
        for face in faces:
            if isolated_vertex not in face:
                continue
            # 找到包含孤立顶点的面并添加连接
            face = list(face)
            face.append(connected_vertex)
            faces.append(face)
        return faces

    def fix_unconnected_edges(self, faces, vertex_connections):
        """
        修复未连接的边，确保网格封闭
        """
        edges = self.get_edges_from_triangles(faces)
        unique_edges = set(tuple(sorted(edge)) for edge in edges)  # 去除重复边

        # 找到未连接的边，进行修复
        for edge in unique_edges:
            if not self.is_edge_connected(edge, vertex_connections):
                # 如果边未连接，添加连接
                self.connect_edge(edge, vertex_connections, faces)

        return faces

    def is_edge_connected(self, edge, vertex_connections):
        """
        检查边是否已经连接
        """
        v1, v2 = edge
        return v2 in vertex_connections.get(v1, []) or v1 in vertex_connections.get(v2, [])

    def connect_edge(self, edge, vertex_connections, faces):
        """
        连接未连接的边
        """
        v1, v2 = edge
        # 将边连接到其他三角形中
        faces = self.add_connection(faces, v1, v2)
        return faces
